﻿// MinNet.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include "Util.hpp"
#include "Dataset.hpp"
#include "Function.hpp"
#include "Layer.hpp"
#include "Optimizer.hpp"

#include <iostream>
#include <string>
#include <time.h>

class ResNet6 :public minnet::model::Model {
public:
    ResNet6() {
        conv1 = minnet::model::Conv2d(6, 1, 7, 3);
        basicblock1 = minnet::model::Sequence(
            minnet::model::Conv2d(6, 6),
            minnet::model::Relu(),
            minnet::model::Conv2d(6, 6)
        );
        basicblock2 = minnet::model::Sequence(
            minnet::model::Conv2d(6, 6),
            minnet::model::Relu(),
            minnet::model::Conv2d(6, 6)
        );
        fc1 = minnet::model::Linear(6 * (3 * 3), 10);
        register_layer(conv1, basicblock1, basicblock2, fc1);
    }
    minnet::Tensor Forward(minnet::Tensor& input) override {
        minnet::Tensor out = conv1(input);
        out = minnet::function::Relu(out);
        out = minnet::function::MaxPool2d(out);
        out = basicblock1(out) + out;
        out = minnet::function::Relu(out);
        out = minnet::function::MaxPool2d(out);
        out = basicblock2(out) + out;
        out = minnet::function::Relu(out);
        out = minnet::function::MaxPool2d(out);
        out = out.reshape(1, -1);
        return fc1(out);
    }
private:
    minnet::model::Conv2d conv1;
    minnet::model::Sequence basicblock1;
    minnet::model::Sequence basicblock2;
    minnet::model::Linear fc1;
};

#define NUM 60000
#define TEST_NUM 10000

#define SHOW_RESULT

int main(int argc, char**argv) {
    if(argc != 2){
        std::cout<<"must input mnist dataset path"<<std::endl;
        return 0;
    }
    
    std::string dataset_root_path = std::string(argv[1]);
    std::string train_data_path = dataset_root_path + "train-images.idx3-ubyte";
    std::string train_label_path = dataset_root_path + "train-labels.idx1-ubyte";
    std::string test_data_path = dataset_root_path + "t10k-images.idx3-ubyte";
    std::string test_label_path = dataset_root_path + "t10k-labels.idx1-ubyte";

    auto src_data = load_mnist(train_data_path, train_label_path);
    auto test_data = load_mnist(test_data_path , test_label_path);
    if (src_data.size() != NUM || test_data.size() != TEST_NUM) {
        std::cout<<"read dataset error!"<<std::endl;
        return 0;
    }
    srand((unsigned)time(NULL));
    std::vector<std::vector<float>> data(NUM);
    std::vector<std::vector<float>> label(NUM);

    std::vector<std::vector<float>> test(TEST_NUM);
    std::vector<std::vector<float>> test_label(TEST_NUM);

    for (int i = 0; i < NUM; i++) {
        data[i] = image_to_vec(src_data[i].second);
        label[i] = std::vector<float>(10, 0.f);
        label[i][src_data[i].first] = 1.f;
    }

    for (int i = 0; i < TEST_NUM; i++) {
        test[i] = image_to_vec(test_data[i].second);
        test_label[i] = std::vector<float>(10, 0.f);
        test_label[i][test_data[i].first] = 1.f;
    }

    ResNet6 net;
    float lr = 0.001;
    minnet::optimizer::SGD opt(net.parameters(), lr, 0.9f);
    minnet::util::Timer timer;
    net.train();
    for (int i = 0; i < 5 ; i++) {
        minnet::shuffle(&data, &label, &src_data);
        float forward_time = 0.f;
        float backward_time = 0.f;
        float total_loss = 0.f;
        float temp_loss = 0.f;
        int count = 0;
        for (int j = 0; j < NUM; j++) {
            minnet::Tensor in;
            in.from_vector_2d(data, j, j + 1);
            minnet::Tensor real;
            real.from_vector_2d(label, j, j + 1);
            minnet::Tensor result = in.reshape(28, 28, 1);
            timer.begin();
            result = net(result);
            real.reshape(1, 10);
            if (minnet::argMax(result) == minnet::argMax(real)) count++;
            minnet::Tensor loss = minnet::function::CrossEntropyLoss(result, real);
            temp_loss += loss.at(0, 0);
            timer.end();
            forward_time += timer.cost();

            timer.begin();
            opt.zero_grad();
            loss.backward();
            opt.step();
            timer.end();
            backward_time += timer.cost();

            if ((j + 1) % 10000 == 0) {
                std::cout<< "epoch: "<< i + 1 <<" train(" << j + 1 << "/" << 60000 << ")" << " loss: " << 
                    temp_loss / (10000.f) <<std::endl;
                total_loss += temp_loss;
                temp_loss = 0.f;
            }
        }
        std::cout << "epoch: " << i + 1 << " forward cost: " << forward_time << " backward cost: " << backward_time <<
            " avg loss: " << total_loss / (NUM * 1.f) << " accuracy in train set: " << count / (NUM * 1.f) << std::endl;
        opt.set_lr(opt.get_lr() / 2.f);
    }

    int count = 0;
    net.eval();
    for (int j = 0; j < TEST_NUM; j++) {
        minnet::Tensor in;
        in.from_vector_2d(test, j, j + 1);
        minnet::Tensor real;
        real.from_vector_2d(test_label, j, j + 1);
        minnet::Tensor result = in.reshape(28, 28, 1);
        result = net(result);
        real.reshape(1, 10);
        if (j < 20) {
            std::cout << "pred: " << minnet::argMax(result) << " real: ";
            std::cout << minnet::argMax(real) << std::endl;
 #ifdef SHOW_RESULT
            cv::Mat temp;
            cv::resize(test_data[j].second, temp, cv::Size(224, 224));
            cv::imshow("result", temp);
            cv::waitKey(0); 
            cv::destroyWindow("result");
#endif // SHOW_RESULT                   
        }
        if (minnet::argMax(result) == minnet::argMax(real)) count++;
    }
    std::cout << "after train accuracy in test set: " << count / (TEST_NUM * 1.f) << std::endl;

    return 0;
}
